import traceback
import cgi
from functools import wraps

class TestFlow(object):
  def __init__(self, name=''):
    self.name = name or 'TestFlow'
    self.tests = []
    self.assertions_count = 0
    self.fail_count = 0
    self.block_fail_count = 0
    self._start = (lambda flow : None)
    self._finish = (lambda flow : None)
    self._setup = (lambda flow : None)
    self._teardown = (lambda flow : None)
  
  def checks(self, checkee=None, *args, **kwargs):
    if checkee:
      assert callable(checkee), "Checked object must be callable"
    
    def _decorator(func):
      
      @wraps(func)
      def _decorated(flow):
        if checkee:
          ret = checkee(*args, **kwargs)
        else:
          ret = None
        return func(flow, ret)
      
      if not _decorated.__doc__:
        _decorated.__doc__ = checkee.__doc__
      self.tests.append(_decorated)
      # return _decorated # Don't return function, no calling outside flow
    
    return _decorator
  
  def start(self, func):
    self._start = func
    
  def finish(self, func):
    self._finish = func
  
  def setup(self, func):
    self._setup = func
  
  def teardown(self, func):
    self._teardown = func
  
  setUp = setup
  set_up = setup
  tearDown = teardown
  tear_down = teardown
    
  def __call__(self, out, up_to=None):
    self.out = out
    up_to = up_to or len(self.tests)
    TestFlow.current = self
    
    self.assertions_count = 0
    self.fail_count = 0
    
    self.msg('Starting tests for %s' % self.name)
    self._start(self)
    
    try:
      for index, test in enumerate(self.tests[:up_to]):
        self.block_fail_count = 0
        self.msg('\nStarting block %s in %s' % (index, test.__module__), b=True)
        if test.__doc__:
          self.msg(test.__doc__, i=True)
        
        self._setup(self)
        try:
          ret = test(self)
          if ret is False or (ret is None and self.block_fail_count):
            raise AssertionError(
              "Test failures prevented this test flow from continuing")
        finally:
          self._teardown(self)
    
    except Exception, x:
      self.error(x)
    
    finally:
      self._finish(self)
      
    self.msg("\nRan %s test(s), %s assertion(s), %s failure(s)\n" % 
      (up_to, self.assertions_count, self.fail_count))  
    
  def msg(self, msg, color=None, i=False, b=False):
    msg = str(msg)
    msg = cgi.escape(msg)
    msg = msg.replace('\n','<br />').replace('\r','<br />')
    if color: msg = ('<span style=\"color:%s\">' % color) + msg + '</span>'
    if i: msg = '<i>' + msg + '</i>'
    if b: msg = '<b>' + msg + '</b>'
    self.out.write(msg + '<br />')
    
  def error(self, exception):
    self.msg('\n' + traceback.format_exc() + '\n' + str(exception), color='red')
    
  def check(self, var, msg='', limit=5):
    self.assertions_count += 1
    assertion_msg = "Assertion %s" % self.assertions_count
    msg = msg or assertion_msg
    if not var:
      self.fail_count += 1
      self.block_fail_count += 1
      self.msg(msg + ' ... Failed', color='red')
      self.msg('----------', color='red')
      self.msg(''.join(traceback.format_stack(limit=limit+1)[:-1]),
        color='red')
      return False
    else:
      self.msg(assertion_msg + ' ... Passed', color='green')
      return True
  
  def check_equal(self, a, b):
    repr_a = repr(a)
    repr_b = repr(b)
    return self.check((a == b), "%s should equal %s" % (repr_a, repr_b))
  
  def check_not_equal(self, a, b):
    repr_a = repr(a)
    repr_b = repr(b)
    return self.check((a != b), "%s should not equal %s" % (repr_a, repr_b))
      
  def check_in(self, a, b):
    repr_a = repr(a)
    repr_b = repr(b)
    return self.check((a in b), "%s should be in %s" % (repr_a, repr_b)) 
  
  def check_not_in(self, a, b):
    repr_a = repr(a)
    repr_b = repr(b)
    return self.check((a not in b), "%s should not be in %s" % (repr_a, repr_b)) 

_default_flow = TestFlow()
import sys
_default_flow.out = sys.stderr
TestFlow.current = TestFlow()

def msg(*args):
  TestFlow.current.msg(*args)

